import jieba
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

import Config
from DJWorkModels.crawler.mysqldata import MyMySQL
from SearchModels.knowledge_base import GetPromptByText


class DJAssistant:
    def __init__(self, LLM):
        self.LLM = LLM
        self.ans_category = ["讲话", "会议", "活动", "考察", "会见", "出访", "函电", "其它", "经济", "政治", "文化",
                             "社会",
                             "生态", "党建", "国防", "无关"]
        self.GetPrompt = GetPromptByText(search_model_name=Config.SEARCH_MODEL_NAME)
        self.ms = MyMySQL(host='innerschool.software-studio.cn', user='dangjian', password='dangjian',
                          database="dangjian")

    def answer(self, ask, history):
        cate_ask = str({
            "prompt":
                "你的一个名为小易的党建助手，请你对下面用户的问题与之前用户的问题进行分析并推测用户的当前问题是否需要从党建数据库中去寻找答案，若你判断问题与党建无关或者不需要使用数据库寻找答案，请不要给出任何分析，直接回答 无关",
            "若你判断问题需要使用数据库寻找答案，请你从以下类别中选择最有可能出现答案的类别，并直接输出类别，不要有任何分析"
            "类别": ["讲话", "会议", "活动", "考察", "会见", "出访", "函电", "其它", "经济", "政治", "文化", "社会",
                     "生态", "国防", "无关"],
            "words": ask
        })
        res = self.LLM.answer_no_stream(cate_ask, history)
        # res是一个字符串，请比较其与self.ans_category中哪一个最匹配
        category = self.determine_category(res, self.ans_category)
        print("category", category)
        if category == "无关":
            Config.DEFAULT_PROMMPT["words"] = ask
            return self.LLM.answer(str(Config.DEFAULT_PROMMPT), ask, history)
        # 如果mysql没有连接则连接
        if not self.ms.is_connected():
            self.ms.connect()
        titles = self.get_article_titles(category)
        # 从各个标题中找到最合适的标题
        titles = self.determine_titles(ask, history, titles)
        # 定位找到的文章
        text = self.get_articles_by_titles(titles)
        # 根据文本去获取关键段并生成提示词
        prompt = self.GetPrompt.get_knowledge_based_answer(text, ask, history)
        # 将提示词输入
        result = self.LLM.answer(prompt, ask, history)
        return result

    def tokenize_chinese(self, text):
        """使用jieba进行中文分词"""
        return jieba.lcut(text)

    def compute_similarity(self, sentence1, sentence2):
        """计算两个句子的相似度"""
        # 创建一个TfidfVectorizer实例，并使用jieba进行分词
        vectorizer = TfidfVectorizer(tokenizer=self.tokenize_chinese, lowercase=False)

        # 将句子转换为TF-IDF特征向量
        tfidf_matrix = vectorizer.fit_transform([sentence1, sentence2])

        # 计算两个句子的余弦相似度
        cosine_sim = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:2])[0][0]

        return cosine_sim

    # 确定输入文本与预定义类别列表中哪个类别的相似度最高
    def determine_category(self, text, categories):
        """
        确定输入文本与预定义类别列表中哪个类别的相似度最高
        """
        highest_similarity = 0
        most_similar_category = None

        # 遍历类别列表，计算每个类别与输入文本的相似度
        for category in categories:
            similarity = self.compute_similarity(text, category)
            if similarity > highest_similarity:
                highest_similarity = similarity
                most_similar_category = category
        if highest_similarity < 0.03:
            return "无关"
        return most_similar_category

    # 找到最靠谱的三个标题
    def determine_titles(self, ask, history, titles):
        """
        确定输入文本与预定义类别列表中哪个类别的相似度最高
        """

        # 遍历titles列表，计算每个标题与输入文本的相似度，并返回前三个最相似的标题
        similarities = []
        # 遍历titles列表，计算每个标题与输入文本的相似度  
        for title in titles:
            similarity = self.compute_similarity(ask, title)
            similarities.append((title, similarity))

            # 根据相似度进行排序，并返回前三个最相似的标题
        similarities.sort(key=lambda x: x[1], reverse=True)
        if len(similarities) < 3:
            # 以列表形式返回title
            return [i[0] for i in similarities]
        else:
            top_three = similarities[:3]
            # 以列表形式返回title
            return [i[0] for i in top_three]

    # 根据种类去数据库中查找对应的文章，首先应该返回标题
    def get_article_titles(self, category):
        if category == "无关":
            return "无关"
        elif category in ["讲话", "会议", "活动", "考察", "会见", "出访", "函电", "其它"]:
            sql = "select 标题 from 文章库 where 类型='%s'" % category
            self.ms.cursor.execute(sql)
            res = self.ms.cursor.fetchall()
            return res
        elif category in ["经济", "政治", "文化", "社会", "生态", "党建", "国防"]:
            sql = "select 标题 from 文章库 where 领域='%s'" % category
            self.ms.cursor.execute(sql)
            res = self.ms.cursor.fetchall()
            # 将res转化为list
            res = [i[0] for i in res]

            return res
        else:
            return "无关"

    # 根据文章标题（一个字符串）去数据库查询并返回文章
    def get_articles_by_titles(self, titles):
        text = ""
        for title in titles:
            text += self.get_article_by_title(title)

        return text

    def get_article_by_title(self, title):
        sql = "select 文章 from 文章库 where 标题='%s'" % title
        self.ms.cursor.execute(sql)
        res = self.ms.cursor.fetchall()
        return res[0][0]
